import json
import time
import logging
import uuid
import random
import string
from typing import Dict, Any, AsyncGenerator

import cloudscraper
from fastapi import HTTPException
from fastapi.responses import StreamingResponse, JSONResponse

from app.core.config import settings
from app.providers.base_provider import BaseProvider
from app.utils.sse_utils import create_sse_data, create_chat_completion_chunk

logger = logging.getLogger(__name__)

class ChatAIGPTProvider(BaseProvider):
    def __init__(self):
        self.scraper = cloudscraper.create_scraper()
        self.base_url = "https://chataigpt.net/wp-admin/admin-ajax.php"
        self.headers = {
            "Accept": "text/event-stream",
            "Accept-Language": "zh-CN,zh;q=0.9,en;q=0.8",
            "Origin": "https://chataigpt.net",
            "Referer": "https://chataigpt.net/",
            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/141.0.0.0 Safari/537.36"
        }

    def _generate_random_string(self, length=5):
        return ''.join(random.choices(string.ascii_lowercase + string.digits, k=length))

    async def chat_completion(self, request_data: Dict[str, Any]) -> StreamingResponse:
        model = request_data.get("model", "gpt-4")
        bot_id = settings.bot_id_map.get(model)
        if not bot_id:
            raise HTTPException(status_code=404, detail=f"模型 '{model}' 不支持。可用模型: {', '.join(settings.known_models)}")

        async def stream_generator() -> AsyncGenerator[bytes, None]:
            request_id = f"chatcmpl-{uuid.uuid4()}"
            try:
                cache_key = self._get_cache_key(request_data, bot_id)

                params = {
                    "action": "aipkit_frontend_chat_stream",
                    "cache_key": cache_key,
                    "bot_id": bot_id,
                    "session_id": str(uuid.uuid4()),
                    "conversation_uuid": str(uuid.uuid4()),
                    "post_id": settings.POST_ID,
                    "_ts": int(time.time() * 1000),
                    "_ajax_nonce": settings.AJAX_NONCE,
                }
                
                # 使用 cloudscraper 的 iter_lines 时，它可能会错误地预解码
                # 我们需要获取原始字节并手动处理
                response = self.scraper.get(self.base_url, params=params, headers=self.headers, stream=True, timeout=settings.API_REQUEST_TIMEOUT)
                response.raise_for_status()
                
                logger.info(f"开始从上游接收流式数据 (bot_id: {bot_id})")
                
                is_error_event = False
                for line in response.iter_lines():
                    if not line:
                        continue
                    
                    # 解码为字符串以进行检查
                    line_str = line.decode('utf-8', errors='ignore')

                    if line_str.startswith("event: error"):
                        is_error_event = True
                        continue

                    if line_str.startswith("data:"):
                        content_str = line_str[len("data:"):].strip()
                        
                        # **关键修复：处理编码问题**
                        # 将可能被错误解码为 latin-1 的字符串还原为原始 UTF-8 字节，再正确解码
                        try:
                            # 尝试直接解析
                            data = json.loads(content_str)
                        except json.JSONDecodeError:
                            # 如果失败，则应用编码修复
                            re_encoded_bytes = content_str.encode('latin-1')
                            correctly_decoded_str = re_encoded_bytes.decode('utf-8')
                            data = json.loads(correctly_decoded_str)

                        # **关键修复：处理上游错误事件**
                        if is_error_event:
                            error_content = data.get("error", "未知的上游错误")
                            logger.error(f"上游服务返回错误: {error_content}")
                            # 在流中发送一个错误块
                            error_chunk = create_chat_completion_chunk(request_id, model, f"Upstream Error: {error_content}", "error")
                            yield create_sse_data(error_chunk)
                            break # 终止流

                        if data.get("finished"):
                            break
                        
                        delta_content = data.get("delta")
                        if delta_content is not None:
                            chunk = create_chat_completion_chunk(request_id, model, delta_content)
                            yield create_sse_data(chunk)

                    elif line_str.startswith("event: done"):
                        break
                
                final_chunk = create_chat_completion_chunk(request_id, model, "", "stop")
                yield create_sse_data(final_chunk)
                yield b"data: [DONE]\n\n"

            except Exception as e:
                logger.error(f"处理流时发生严重错误: {e}", exc_info=True)
                error_message = f"内部服务器错误: {str(e)}"
                error_chunk = create_chat_completion_chunk(request_id, model, error_message, "error")
                yield create_sse_data(error_chunk)
                yield b"data: [DONE]\n\n"

        return StreamingResponse(stream_generator(), media_type="text/event-stream")

    def _get_cache_key(self, request_data: Dict[str, Any], bot_id: str) -> str:
        messages = request_data.get("messages", [])
        last_message = next((m['content'] for m in reversed(messages) if m['role'] == 'user'), "")

        form_data = {
            "action": (None, "aipkit_cache_sse_message"),
            "message": (None, last_message),
            "_ajax_nonce": (None, settings.AJAX_NONCE),
            "bot_id": (None, bot_id),
            "user_client_message_id": (None, f"aipkit-client-msg-{bot_id}-{int(time.time() * 1000)}-{self._generate_random_string()}"),
        }
        
        try:
            # 移除 Accept header，因为此请求不返回 event-stream
            post_headers = self.headers.copy()
            post_headers["Accept"] = "*/*"
            response = self.scraper.post(self.base_url, headers=post_headers, files=form_data, timeout=settings.API_REQUEST_TIMEOUT)
            response.raise_for_status()
            data = response.json()

            if data.get("success") and data.get("data", {}).get("cache_key"):
                cache_key = data["data"]["cache_key"]
                logger.info(f"成功获取 cache_key: {cache_key}")
                return cache_key
            else:
                error_msg = data.get('data', {}).get('message', '未知错误')
                logger.error(f"获取 cache_key 失败，响应: {error_msg}")
                raise HTTPException(status_code=500, detail=f"无法从上游获取 cache_key: {error_msg}")
        except Exception as e:
            logger.error(f"请求 cache_key 时发生异常: {e}", exc_info=True)
            raise HTTPException(status_code=502, detail=f"请求上游 cache_key 失败: {str(e)}")

    async def get_models(self) -> JSONResponse:
        model_data = {
            "object": "list",
            "data": [
                {"id": name, "object": "model", "created": int(time.time()), "owned_by": "lzA6"}
                for name in settings.known_models
            ]
        }
        return JSONResponse(content=model_data)
